/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.translation.functions;



import com.bff.gaia.unified.runners.core.DoFnRunner;

import com.bff.gaia.unified.runners.core.DoFnRunners;

import com.bff.gaia.unified.runners.core.construction.SerializablePipelineOptions;

import com.bff.gaia.unified.runners.gaia.GaiaPipelineOptions;

import com.bff.gaia.unified.runners.gaia.metrics.DoFnRunnerWithMetricsUpdate;

import com.bff.gaia.unified.runners.gaia.translation.utils.GaiaClassloading;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.DoFnSchemaInformation;

import com.bff.gaia.unified.sdk.transforms.join.RawUnionValue;

import com.bff.gaia.unified.sdk.transforms.reflect.DoFnInvoker;

import com.bff.gaia.unified.sdk.transforms.reflect.DoFnInvokers;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.sdk.values.PCollectionView;

import com.bff.gaia.unified.sdk.values.TupleTag;

import com.bff.gaia.unified.sdk.values.WindowingStrategy;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Lists;

import com.bff.gaia.api.common.functions.RichMapPartitionFunction;

import com.bff.gaia.api.common.functions.RuntimeContext;

import com.bff.gaia.configuration.Configuration;

import com.bff.gaia.util.Collector;



import java.util.List;

import java.util.Map;

import java.util.Optional;



/**

 * Encapsulates a {@link DoFn} inside a Gaia {@link

 * com.bff.gaia.api.common.functions.RichMapPartitionFunction}.

 *

 * <p>We get a mapping from {@link TupleTag} to output index and must tag

 * all outputs with the output number. Afterwards a filter will filter out those elements that are

 * not to be in a specific output.

 */

public class GaiaDoFnFunction<InputT, OutputT>

    extends RichMapPartitionFunction<WindowedValue<InputT>, WindowedValue<OutputT>> {



  private final SerializablePipelineOptions serializedOptions;



  private final DoFn<InputT, OutputT> doFn;

  private final String stepName;

  private final Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputs;



  private final WindowingStrategy<?, ?> windowingStrategy;



  private final Map<TupleTag<?>, Integer> outputMap;

  private final TupleTag<OutputT> mainOutputTag;

  private final Coder<InputT> inputCoder;

  private final Map<TupleTag<?>, Coder<?>> outputCoderMap;

  private final DoFnSchemaInformation doFnSchemaInformation;



  private transient DoFnInvoker<InputT, OutputT> doFnInvoker;



  public GaiaDoFnFunction(

      DoFn<InputT, OutputT> doFn,

      String stepName,

      WindowingStrategy<?, ?> windowingStrategy,

      Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputs,

      PipelineOptions options,

      Map<TupleTag<?>, Integer> outputMap,

      TupleTag<OutputT> mainOutputTag,

      Coder<InputT> inputCoder,

      Map<TupleTag<?>, Coder<?>> outputCoderMap,

      DoFnSchemaInformation doFnSchemaInformation) {



    this.doFn = doFn;

    this.stepName = stepName;

    this.sideInputs = sideInputs;

    this.serializedOptions = new SerializablePipelineOptions(options);

    this.windowingStrategy = windowingStrategy;

    this.outputMap = outputMap;

    this.mainOutputTag = mainOutputTag;

    this.inputCoder = inputCoder;

    this.outputCoderMap = outputCoderMap;

    this.doFnSchemaInformation = doFnSchemaInformation;

  }



  @Override

  public void mapPartition(

      Iterable<WindowedValue<InputT>> values, Collector<WindowedValue<OutputT>> out)

      throws Exception {



    RuntimeContext runtimeContext = getRuntimeContext();



    DoFnRunners.OutputManager outputManager;

    if (outputMap.size() == 1) {

      outputManager = new GaiaDoFnFunction.DoFnOutputManager(out);

    } else {

      // it has some additional outputs

      outputManager = new GaiaDoFnFunction.MultiDoFnOutputManager((Collector) out, outputMap);

    }



    List<TupleTag<?>> additionalOutputTags = Lists.newArrayList(outputMap.keySet());



    DoFnRunner<InputT, OutputT> doFnRunner =

        DoFnRunners.simpleRunner(

            serializedOptions.get(),

            doFn,

            new GaiaSideInputReader(sideInputs, runtimeContext),

            outputManager,

            mainOutputTag,

            additionalOutputTags,

            new GaiaNoOpStepContext(),

            inputCoder,

            outputCoderMap,

            windowingStrategy,

            doFnSchemaInformation);



    if ((serializedOptions.get().as(GaiaPipelineOptions.class)).getEnableMetrics()) {

      doFnRunner = new DoFnRunnerWithMetricsUpdate<>(stepName, doFnRunner, getRuntimeContext());

    }



    doFnRunner.startBundle();



    for (WindowedValue<InputT> value : values) {

      doFnRunner.processElement(value);

    }



    doFnRunner.finishBundle();

  }



  @Override

  public void open(Configuration parameters) throws Exception {

    doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn);

  }



  @Override

  public void close() throws Exception {

    try {

      Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);

    } finally {

      GaiaClassloading.deleteStaticCaches();

    }

  }



  static class DoFnOutputManager implements DoFnRunners.OutputManager {



    private Collector collector;



    DoFnOutputManager(Collector collector) {

      this.collector = collector;

    }



    @Override

    @SuppressWarnings("unchecked")

    public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {

      collector.collect(

          WindowedValue.of(

              new RawUnionValue(0 /* single output */, output.getValue()),

              output.getTimestamp(),

              output.getWindows(),

              output.getPane()));

    }

  }



  static class MultiDoFnOutputManager implements DoFnRunners.OutputManager {



    private Collector<WindowedValue<RawUnionValue>> collector;

    private Map<TupleTag<?>, Integer> outputMap;



    MultiDoFnOutputManager(

        Collector<WindowedValue<RawUnionValue>> collector, Map<TupleTag<?>, Integer> outputMap) {

      this.collector = collector;

      this.outputMap = outputMap;

    }



    @Override

    public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {

      collector.collect(

          WindowedValue.of(

              new RawUnionValue(outputMap.get(tag), output.getValue()),

              output.getTimestamp(),

              output.getWindows(),

              output.getPane()));

    }

  }

}